package com.unicdata.ai.demo_2af.controller;

import com.baomidou.mybatisplus.extension.api.R;
import com.unicdata.ai.demo_2af.common.entity.User;
import com.unicdata.ai.demo_2af.common.result.Result;
import com.unicdata.ai.demo_2af.service.TwoFactorService;
import org.apache.shiro.SecurityUtils;
import org.apache.shiro.subject.Subject;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.Map;

/**
 * 2FA 控制器
 *
 * @author lhn
 * @date 2025/06/26 16:44
 * @since 1.0.0
 **/
@RestController
@RequestMapping("/2fa")
public class TwoFactorController {

    @Autowired
    private TwoFactorService twoFactorService;

    /**
     * 生成2FA密钥和二维码URL
     */
    @PostMapping("/setup")
    public Result <Map<String, String>> setup() {
        try {
            Subject subject = SecurityUtils.getSubject();
            User user = (User) subject.getPrincipal();

            if (user.getIsTwoFactorEnabled() != null && user.getIsTwoFactorEnabled() == 1) {
                return Result.error("2FA已经启用");
            }

            Map<String, String> twoFactorData = twoFactorService.generateTwoFactorSecret(user.getUsername());
            return Result.success(twoFactorData);
        } catch (Exception e) {
            return Result.error("生成2FA密钥失败：" + e.getMessage());
        }
    }

    /**
     * 验证并启用2FA
     */
    @PostMapping("/enable")
    public Result<String> enable(@RequestParam("code") int code, @RequestParam("secret") String secret) {
        try {
            Subject subject = SecurityUtils.getSubject();
            User user = (User) subject.getPrincipal();

            if (twoFactorService.validateCode(secret, code)) {
                twoFactorService.enableTwoFactor(user.getId(), secret);
                return Result.success("2FA启用成功");
            }
            return Result.error("验证码错误");
        } catch (Exception e) {
            return Result.error("启用2FA失败：" + e.getMessage());
        }
    }

    /**
     * 验证2FA码
     */
    @PostMapping("/verify")
    public Result<String> verify(@RequestParam("code") int code) {
        try {
            Subject subject = SecurityUtils.getSubject();
            User user = (User) subject.getPrincipal();

            if (user.getIsTwoFactorEnabled() != 1) {
                return Result.success("未启用2FA");
            }

            if (twoFactorService.validateCode(user.getTwoFactorSecret(), code)) {
                return Result.success("验证成功");
            }

            // 验证失败，注销当前会话
            subject.logout();
            return Result.error("验证码错误");
        } catch (Exception e) {
            return Result.error("验证2FA失败：" + e.getMessage());
        }
    }
}